#!/usr/bin/env python
# # -*- coding: utf-8 -*-

"""
@File:      DateCheckTool.py
@Author:    Jim.Dai.Cn
@Date:      2024/6/11 19:00
@Desc:         
"""
from datetime import datetime

import pytz
from qwen_agent.tools.base import BaseTool, register_tool
import json5
import urllib.parse
import ast
import copy
from qwen_agent.agents import Assistant
from ..config.config import *
from ..utils.global_functions import *

'''
    时间矛盾，投标文件递交截止时间2024-5-23 但是 开标时间 2024-5-22
    写定关键业务流程：招投开评定
    寻找流程节点中所有的日期，包括潜在的日期冲突
    转换为时间戳，进行比较
    todo：不同招标流程，按照招投标法规，时间限定不同。
'''


@register_tool('date_check')
class DateCheckTool(BaseTool):
    # The `description` tells the agent the functionality of this tool.
    description = 'Check date and time errors in a doc file with given path.'
    # The `parameters` tell the agent what input parameters the tool has.
    parameters = [
        {'name': 'path_doc', 'type': 'string', 'description': 'The full of path of doc used for analyse', 'required': True}
    ]
    # 业务规则：依法必招项目 [公告发布时间] 到 [投标文件递交截止时间]，不得少于20日
    array_baseline_questions = [
        '文字最开头的撰写日期',
        '文字末尾的完稿日期',
        '公告发布日期',
        '招标文件发售开始日期',
        '招标文件发售截止时间',
        '投标截止时间',
        '开标时间',
        '评标时间',
        '中标候选人公示时间',
        '中标通知时间',
        '合同签订时间'
    ]
    array_other_date_check = []
    str_baseline = ''
    array_chat_history_baseline = []
    _instance = None
    assistant = None
    _doc_content = None
    rules = [
        # '上文开头部分或标题中的日期应与文末的日期一致',
        # '形如"自/从 A年月日 至/到 B年月日"定义的日期范围中，"B年月日"应早于"A年月日"1天',
        # '形如"A年月日 至/到 B年月日"定义的日期范围应大于1天',
        # '形如"自/从 {start_date} 至/到 {end_date}"定义的日期范围中 {end_date} 应不小于 {start_date} 1天',
        # '{start_date}到{end_date}的时间范围中 {start_date} 应早于 {end_date}',
        # '形如"A\s*年B\s*月C\s*日 至/到 X\s*年Y\s*月Z\s*日"表述，"A\s*年B\s*月C\s*日"应早于"X\s*年Y\s*月Z\s*日"',
        # '形如"A\s*年B\s*月C\s*日 至/到 截止时间"的表述，首先推断截止时间的具体日期，再检查 A\s*年B\s*月C\s*日 应早于对应时间',
        # '以截止时间定义的时间范围，首先推断截止时间的具体日期，再检查开始时间应早于结束时间',
        # '以当前时间定义的时间范围，首先推断当前时间的具体日期，再检查开始时间应早于结束时间',
        # '日期应早于合同期限允许的最大期限',
        '撰写日期应早于或等于发布日期',
        '日期范围的开始日期应早于结束日期',
        '如果两个日期的月份与日相同而年份不同，检查是否是潜在的矛盾或错误',
        '日期范围缺少开始或结束的，尝试推断具体日期，并检查开始时间应早于结束时间',
        '项目名称中的年份不应与文本中实际实施年份矛盾'
        # '按照发布公告、下载文件、递交文件、开评标、合同签订、履约完毕、合同期满的顺序，相关日期先后次序应合理'
    ]

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super(DateCheckTool, cls).__new__(cls)
            cls._instance.__init__(*args, **kwargs)
        return cls._instance

    def __init__(self, *args, **kwargs):
        super(DateCheckTool, self).__init__(*args, **kwargs)

    def call(self, params: str, **kwargs) -> str:
        # `params` are the arguments generated by the LLM agent.
        str_path = json5.loads(params)['path_doc']
        self.assistant = Assistant(llm=llm_cfg_llama_cpp_qwen2_7B,
                                   # system_message='你是一位专业的招标代理，请对照规则检查用户提供的文本，严格按照提示回答问题。请勿篡改原文，只用原文回答，言简意赅。',
                                   system_message='你是一位精通Python的招标代理，请对照规则检查用户提供的文本，勿篡改原文，严格以Python列表方式回答问题，直接回答原文内容的数组，不需要任何解释。',
                                   function_list=['date_range_compare_duration', 'date_equal', 'date_earlier_than', 'code_interpreter'])
                                    # files=['DataWizards/tools/rules/date.txt']
                                   # files=[str_path])
        self._doc_content = read_from_doc_to_string(str_path)
        # print(self._doc_content)
        list_ret = self.find_date_contradiction()
        return list_ret
        # return self.date_to_ts(prompt)

    def find_date_contradiction(self) -> list:
        # text_2026 = self.llm_get_answer(f"请列出所有与2026年有关的表述。")
        # print(text_2026)
        # return
        # self.array_chat_history_baseline = []
        self.array_chat_history_baseline = [{'role': 'user', 'content': self._doc_content}]
        flag_create_baseline = False
        if flag_create_baseline:
            self.str_baseline = ""
            for pre_q in self.array_baseline_questions:
                self.llm_prepare_baseline(pre_q)
            # print(f'        上下文准备完毕。{self.array_chat_history_baseline}')
            print(f'        baseline str: {self.str_baseline}')
            if self.str_baseline != "":
                self.array_chat_history_baseline = [{'role': 'user', 'content': self._doc_content+self.str_baseline}]
        arr_result = []
        rule_idx = 1
        for rule in self.rules:
            print(f"        规则检查 #{rule_idx} ：{rule}")
            answer = self.llm_get_answer(f"根据规则：{rule}，深入检查文中日期相关的矛盾或错误，不要修改文本内容，可能有多个错误，不要遗漏任何一处。请以列表方式列结果，错误的以字符串列表方式列引起错误的原文，无错误返回空列表。只给结果不需要解释。")
            print(f"        规则检查 #{rule_idx} ：=> {answer}")
            rule_idx = rule_idx + 1
            if (answer != '[]') and (answer is not None) and (answer != []):
                try:
                    ast_ans = ast.literal_eval(answer)
                    arr_result.append(ast_ans)
                except Exception as e:
                    print(f" => exception {e}")
        # history.append({'role': 'user','content': '请根据规则：日期范围的开始日期应该早于结束日期，深入检查文本中多个日期范围的错误及矛盾，不要遗漏任何一处错误，请用数组方式列出全部结果，只写原文,不需要其他内容，不要忽略空格或者转义。'})
        # history = [{'role': 'user', 'content': '请根据规则：日期及时间范围的开始时间应该小于结束时间，检查文本中时间日期范围的错误及矛盾，请用数组方式列出全部结果，只写原文,不需要其他内容，不要忽略空格或者转义。'}]
        # history = []
        # history = doc_to_history(arr_doc,'请以数组方式列出下述文本中时间日期范围的错误及矛盾，请列出全部，只写原文,不需要其他内容，不要忽略空格或者转义。')
        # history.insert(0, {'role': 'system', 'content': '假设当前日期为文件开头部分所写的文件发布日期当天。日期范围形式：自某年月日至某年月日；日期形式：某年月日。'})
        # history.append({'role': 'user', 'content': '请列出上述内容中所有存在的时间及日期矛盾，只写原文。'})
        # history.append({'role': 'user', 'content': '请以数组方式列出上述文本中时间日期范围的错误及矛盾，请列出全部，只写原文,不需要其他内容。'})
        # history = [{'role': 'user', 'content': '请以数组方式列出文本中时间日期范围的错误及矛盾，请列出全部，只写原文,不需要其他内容，不要忽略空格或者转义。'}]
        # history = [{'role': 'user', 'content': '请根据规则：文件开始与结尾签署时间应该一致，检查文本中时间日期范围的错误及矛盾，请用数组方式列出全部结果，只写原文,不需要其他内容，不要忽略空格或者转义。'}]
        # 。如果工具检查结果是"正常"则跳过继续检查下一处错误。
        # history.append({'role': 'user', 'content': '如存在，请照抄原文矛盾部分。如不存在，请只回答：无。'})
        # history = [{'role': 'user', 'content': '请列出上述内容中所有时间及日期的矛盾原文。'}]
        print(f"      date scanner done, result = {arr_result}")
        # arr_return = ast.literal_eval(arr_result)
        # print(f" => {arr_return}")
        # return arr_return
        return arr_result

    def llm_get_answer(self, question):
        user_text = {'role': 'user', 'content': question}
        # sys_text = {'role': 'system', 'content': self.str_baseline}
        # tmp_his = [sys_text, user_text]
        tmp_his = copy.deepcopy(self.array_chat_history_baseline)
        # tmp_his = [user_text]
        tmp_his.append(user_text)
        reply = []
        for response in self.assistant.run(messages=tmp_his):
            reply.extend(response)
        return reply[-1]['content']

    def llm_prepare_baseline(self, question):
        user_text = {'role': 'user', 'content': self._doc_content + "根据上文回答："+question+'是什么？只用原文回答，40字以内。原文未给出请回答"未知"，不需要解释。'}
        print(f"        baseline: {question}", end=' ')
        reply = []
        for response in self.assistant.run(messages=[user_text]):
            reply.extend(response)
        b_asw = reply[-1]['content']
        print(f" => {b_asw}", end=" ")
        if (b_asw != '未知') and (b_asw is not None) and (len(b_asw) <= 40):
            # self.array_chat_history_baseline.append(user_text)
            # self.array_chat_history_baseline.append({'role': 'assistant', 'content': reply[-1]['content']})
            self.str_baseline = self.str_baseline + "。" + question + b_asw
            print("Saved to baseline")
        else:
            print("未知，不予更新到baseline。")

    def date_to_ts(self, params: str) -> str:
        prompt = params + "转换为YYYY-MM-DD HH:MM:SS格式是多少？请直接给出具体结果。请只回答答案，例如：问题是文件名是什么？回答：文件全称，而不要回答：文件名是文件全称。"
        ts_str = self.llm_get_answer(prompt)
        tz_info = pytz.timezone('Asia/Shanghai')  # 北京时间
        # todo : 时区设置修正
        dt_obj = datetime.strptime(ts_str, "%Y-%m-%d %H:%M:%S")
        dt_obj = tz_info.localize(dt_obj)
        ts = dt_obj.timestamp()
        print(ts)
        return ts


@register_tool('date_equal')
class DateEqual(BaseTool):
    description = 'Check if {first_date} is equal to {second_date}. Return "错误" if not equal; return "正常" if equal.'
    # The `parameters` tell the agent what input parameters the tool has.
    parameters = [
        # {'name': 'origin_text', 'type': 'string', 'description': 'the origin text which contains date or date range','required': True},
        {'name': 'first_date', 'type': 'datetime', 'description': 'the first date', 'required': True},
        {'name': 'second_date', 'type': 'datetime', 'description': 'the second date', ' required': True}
    ]

    def call(self, params: str, **kwargs) -> str:
        # `params` are the arguments generated by the LLM agent.
        # str_origin = json5.loads(params)['origin_text']
        str_start = json5.loads(params)['first_date']
        str_end = json5.loads(params)['second_date']
        if str_start == '' or str_start is None:
            return "正常"
        if str_end == '' or str_end is None:
            return "正常"
        # print(f"date range check => {str_start} | {str_end} | {int_duration}")
        print(f"          date equal check => {str_start} | {str_end}", end=" ")
        result = self.compare_dates(str_start, str_end)
        print(f" result = {result}")
        return result
        # return self.date_to_ts(prompt)

    def compare_dates(self, str_date_start, str_date_end):
        int_days_between = days_between(str_date_start, str_date_end, '%Y-%m-%d')
        print(f"      compare dates: ({str_date_start}, {str_date_end}]", end=" ")
        if int_days_between == 0:
            return '正常'
        else:
            return '错误'


@register_tool('date_earlier_than')
class DateEarlierThan(BaseTool):
    description = 'Check if {first_date} is earlier than {second_date}. Return "错误" if late or equal; return "正常" if earlier.'
    # The `parameters` tell the agent what input parameters the tool has.
    parameters = [
        # {'name': 'origin_text', 'type': 'string', 'description': 'the origin text which contains date or date range','required': True},
        {'name': 'first_date', 'type': 'datetime', 'description': 'the first date', 'required': True},
        {'name': 'second_date', 'type': 'datetime', 'description': 'the second date', ' required': True}
    ]

    def call(self, params: str, **kwargs) -> str:
        # `params` are the arguments generated by the LLM agent.
        # str_origin = json5.loads(params)['origin_text']
        str_start = json5.loads(params)['first_date']
        str_end = json5.loads(params)['second_date']
        if str_start == '' or str_start is None:
            return "正常"
        if str_end == '' or str_end is None:
            return "正常"
        # print(f"date range check => {str_start} | {str_end} | {int_duration}")
        print(f"          date late check => {str_start} | {str_end}", end=" ")
        # date_S = datetime.strptime(str_start, '%Y年%m月%d日')
        # date_E = datetime.strptime(str_end, '%Y年%m月%d日')
        # result = self.compare_dates(date_S, date_E)
        result = self.compare_dates(str_start, str_end)
        print(f" result = {result}")
        return result
        # return self.date_to_ts(prompt)

    def compare_dates(self, str_date_start, str_date_end):
        int_days_between = days_between(str_date_start, str_date_end, '%Y-%m-%d')
        print(f"      compare dates: ({str_date_start}, {str_date_end}]", end=" ")
        if int_days_between <= 0:
            return '错误'
        else:
            return '正常'


@register_tool('date_range_compare_duration')
class DateRangeCompareDuration(BaseTool):
    description = 'Check if the duration of date range from {date_start}(excluded) to {date_end}(included) is NOT LESS than given {duration}. e.g. 2021年1月1日至2021年1月5日 不小于 4天 => 正常。 Return "错误" if contradiction; return "正常" if normal.'
    # The `parameters` tell the agent what input parameters the tool has.
    parameters = [
        # {'name': 'origin_text', 'type': 'string', 'description': 'the origin text which contains date or date range','required': True},
        {'name': 'date_start', 'type': 'datetime', 'description': 'the date start, excluded from range', 'required': True},
        {'name': 'date_end', 'type': 'datetime', 'description': 'the date end, included in range', ' required': True},
        {'name': 'duration', 'type': 'int', 'description': 'The duration in days', 'required': True}
    ]

    def call(self, params: str, **kwargs) -> str:
        # `params` are the arguments generated by the LLM agent.
        # str_origin = json5.loads(params)['origin_text']
        str_start = json5.loads(params)['date_start']
        str_end = json5.loads(params)['date_end']
        int_duration = int(json5.loads(params)['duration'])  # 日 工作日
        if str_start == '' or str_start is None:
            return "正常"
        if str_end == '' or str_end is None:
            return "正常"
        # print(f"date range check => {str_start} | {str_end} | {int_duration}")
        print(f"          date range check => {str_start} | {str_end} | {int_duration}", end=" ")
        result = self.compare_dates(str_start, str_end, int_duration)
        print(f" result = {result}")
        return result
        # return self.date_to_ts(prompt)

    def compare_dates(self, str_date_start, str_date_end, int_days_compare):
        int_days_between = days_between(str_date_start, str_date_end, '%Y-%m-%d %H:%M:%S')
        print(f"      compare dates: ({str_date_start}, {str_date_end}] vs {int_days_compare}", end=" ")
        if int_days_between < 0:
            return '错误'
        if int_days_between >= int_days_compare:
            return '正常'
        else:
            return '错误'

    def comapre_business_days(self, str_date_start, str_date_end, int_days_compare):
        int_days_between = business_days_between(str_date_start, str_date_end, '%Y-%m-%d %H:%M:%S')
        return int_days_between >= int_days_compare
